/*
 * Copyright 2023 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

@file:OptIn(ExperimentalSerializationApi::class)

package com.google.firebase.ai.type

import com.google.firebase.ai.common.util.FirstOrdinalSerializer
import java.util.Calendar
import kotlinx.serialization.ExperimentalSerializationApi
import kotlinx.serialization.KSerializer
import kotlinx.serialization.SerialName
import kotlinx.serialization.Serializable
import kotlinx.serialization.json.JsonNames

/**
 * A `Candidate` represents a single response generated by the model for a given request.
 *
 * @property content The actual content generated by the model.
 * @property safetyRatings A list of [SafetyRating]s describing the generated content.
 * @property citationMetadata Metadata about the sources used to generate this content.
 * @property finishReason The reason the model stopped generating content, if it exist.
 * @property groundingMetadata Metadata returned to the client when the model grounds its response.
 * @property urlContextMetadata Metadata returned to the client when the [UrlContext] tool is
 * enabled.
 */
public class Candidate
@OptIn(PublicPreviewAPI::class)
internal constructor(
  public val content: Content,
  public val safetyRatings: List<SafetyRating>,
  public val citationMetadata: CitationMetadata?,
  public val finishReason: FinishReason?,
  public val groundingMetadata: GroundingMetadata?,
  @property:PublicPreviewAPI public val urlContextMetadata: UrlContextMetadata?
) {

  @OptIn(PublicPreviewAPI::class)
  @Serializable
  internal data class Internal(
    val content: Content.Internal? = null,
    val finishReason: FinishReason.Internal? = null,
    val safetyRatings: List<SafetyRating.Internal>? = null,
    val citationMetadata: CitationMetadata.Internal? = null,
    val groundingMetadata: GroundingMetadata.Internal? = null,
    val urlContextMetadata: UrlContextMetadata.Internal? = null
  ) {

    @OptIn(PublicPreviewAPI::class)
    internal fun toPublic(): Candidate {
      val safetyRatings = safetyRatings?.mapNotNull { it.toPublic() }.orEmpty()
      val citations = citationMetadata?.toPublic()
      val finishReason = finishReason?.toPublic()
      val groundingMetadata = groundingMetadata?.toPublic()
      val urlContextMetadata = urlContextMetadata?.toPublic()

      return Candidate(
        this.content?.toPublic() ?: content("model") {},
        safetyRatings,
        citations,
        finishReason,
        groundingMetadata,
        urlContextMetadata
      )
    }
  }
}

/**
 * An assessment of the potential harm of some generated content.
 *
 * The rating will be restricted to a particular [category].
 *
 * @property category The category of harm being assessed (for example, Hate speech).
 * @property probability The likelihood of the content causing harm.
 * @property probabilityScore A numerical score representing the probability of harm, between `0`
 * and `1`.
 * @property blocked Indicates whether the content was blocked due to safety concerns.
 * @property severity The severity of the potential harm.
 * @property severityScore A numerical score representing the severity of harm.
 */
public class SafetyRating
internal constructor(
  public val category: HarmCategory,
  public val probability: HarmProbability,
  public val probabilityScore: Float = 0f,
  public val blocked: Boolean? = null,
  public val severity: HarmSeverity? = null,
  public val severityScore: Float? = null
) {

  @Serializable
  internal data class Internal
  @JvmOverloads
  constructor(
    val category: HarmCategory.Internal? = null,
    val probability: HarmProbability.Internal? = null,
    val blocked: Boolean? = null, // TODO(): any reason not to default to false?
    val probabilityScore: Float? = null,
    val severity: HarmSeverity.Internal? = null,
    val severityScore: Float? = null,
  ) {

    internal fun toPublic() =
      // Due to a bug in the backend, it's possible that we receive
      // an invalid `SafetyRating` value, without either category or
      // probability. We return null in those cases to enable
      // filtering by the higher level types.
      if (category == null || probability == null) {
        null
      } else {
        SafetyRating(
          category = category.toPublic(),
          probability = probability.toPublic(),
          probabilityScore = probabilityScore ?: 0f,
          blocked = blocked,
          severity = severity?.toPublic(),
          severityScore = severityScore
        )
      }
  }
}

/**
 * A collection of source attributions for a piece of content.
 *
 * @property citations A list of individual cited sources and the parts of the content to which they
 * apply.
 */
public class CitationMetadata internal constructor(public val citations: List<Citation>) {

  @Serializable
  internal data class Internal
  @OptIn(ExperimentalSerializationApi::class)
  internal constructor(@JsonNames("citations") val citationSources: List<Citation.Internal>) {

    internal fun toPublic() = CitationMetadata(citationSources.map { it.toPublic() })
  }
}

/**
 * Represents a citation of content from an external source within the model's output.
 *
 * When the language model generates text that includes content from another source, it should
 * provide a citation to properly attribute the original source. This class encapsulates the
 * metadata associated with that citation.
 *
 * @property title The title of the cited source, if available.
 * @property startIndex The (inclusive) starting index within the model output where the cited
 * content begins.
 * @property endIndex The (exclusive) ending index within the model output where the cited content
 * ends.
 * @property uri The URI of the cited source, if available.
 * @property license The license under which the cited content is distributed under, if available.
 * @property publicationDate The date of publication of the cited source, if available.
 */
public class Citation
internal constructor(
  public val title: String? = null,
  public val startIndex: Int = 0,
  public val endIndex: Int,
  public val uri: String? = null,
  public val license: String? = null,
  public val publicationDate: Calendar? = null
) {

  @Serializable
  internal data class Internal(
    val title: String? = null,
    val startIndex: Int = 0,
    val endIndex: Int,
    val uri: String? = null,
    val license: String? = null,
    val publicationDate: Date? = null,
  ) {

    internal fun toPublic(): Citation {
      val publicationDateAsCalendar =
        publicationDate?.let {
          val calendar = Calendar.getInstance()
          // Internal `Date.year` uses 0 to represent not specified. We use 1 as default.
          val year = if (it.year == null || it.year < 1) 1 else it.year
          // Internal `Date.month` uses 0 to represent not specified, or is 1-12 as months. The
          // month as
          // expected by [Calendar] is 0-based, so we subtract 1 or use 0 as default.
          val month = if (it.month == null || it.month < 1) 0 else it.month - 1
          // Internal `Date.day` uses 0 to represent not specified. We use 1 as default.
          val day = if (it.day == null || it.day < 1) 1 else it.day
          calendar.set(year, month, day)
          calendar
        }
      return Citation(
        title = title,
        startIndex = startIndex,
        endIndex = endIndex,
        uri = uri,
        license = license,
        publicationDate = publicationDateAsCalendar
      )
    }

    @Serializable
    internal data class Date(
      /** Year of the date. Must be between 1 and 9999, or 0 for no year. */
      val year: Int? = null,
      /** 1-based index for month. Must be from 1 to 12, or 0 to specify a year without a month. */
      val month: Int? = null,
      /**
       * Day of a month. Must be from 1 to 31 and valid for the year and month, or 0 to specify a
       * year by itself or a year and month where the day isn't significant.
       */
      val day: Int? = null,
    )
  }
}

/**
 * Represents the reason why the model stopped generating content.
 *
 * @property name The name of the finish reason.
 * @property ordinal The ordinal value of the finish reason.
 */
public class FinishReason private constructor(public val name: String, public val ordinal: Int) {

  @Serializable(Internal.Serializer::class)
  internal enum class Internal {
    UNKNOWN,
    @SerialName("FINISH_REASON_UNSPECIFIED") UNSPECIFIED,
    STOP,
    MAX_TOKENS,
    SAFETY,
    RECITATION,
    OTHER,
    BLOCKLIST,
    PROHIBITED_CONTENT,
    SPII,
    MALFORMED_FUNCTION_CALL;

    internal object Serializer : KSerializer<Internal> by FirstOrdinalSerializer(Internal::class)

    internal fun toPublic() =
      when (this) {
        MAX_TOKENS -> FinishReason.MAX_TOKENS
        RECITATION -> FinishReason.RECITATION
        SAFETY -> FinishReason.SAFETY
        STOP -> FinishReason.STOP
        OTHER -> FinishReason.OTHER
        BLOCKLIST -> FinishReason.BLOCKLIST
        PROHIBITED_CONTENT -> FinishReason.PROHIBITED_CONTENT
        SPII -> FinishReason.SPII
        MALFORMED_FUNCTION_CALL -> FinishReason.MALFORMED_FUNCTION_CALL
        else -> FinishReason.UNKNOWN
      }
  }
  public companion object {
    /** A new and not yet supported value. */
    @JvmField public val UNKNOWN: FinishReason = FinishReason("UNKNOWN", 0)

    /** Model finished successfully and stopped. */
    @JvmField public val STOP: FinishReason = FinishReason("STOP", 1)

    /** Model hit the token limit. */
    @JvmField public val MAX_TOKENS: FinishReason = FinishReason("MAX_TOKENS", 2)

    /** [SafetySetting] prevented the model from outputting content. */
    @JvmField public val SAFETY: FinishReason = FinishReason("SAFETY", 3)

    /**
     * The token generation was stopped because the response was flagged for unauthorized citations.
     */
    @JvmField public val RECITATION: FinishReason = FinishReason("RECITATION", 4)

    /** Model stopped for another reason. */
    @JvmField public val OTHER: FinishReason = FinishReason("OTHER", 5)

    /** Token generation stopped because the content contains forbidden terms. */
    @JvmField public val BLOCKLIST: FinishReason = FinishReason("BLOCKLIST", 6)

    /** Token generation stopped for potentially containing prohibited content. */
    @JvmField public val PROHIBITED_CONTENT: FinishReason = FinishReason("PROHIBITED_CONTENT", 7)

    /**
     * Token generation stopped because the content potentially contains Sensitive Personally
     * Identifiable Information (SPII).
     */
    @JvmField public val SPII: FinishReason = FinishReason("SPII", 8)

    /** The function call generated by the model is invalid. */
    @JvmField
    public val MALFORMED_FUNCTION_CALL: FinishReason = FinishReason("MALFORMED_FUNCTION_CALL", 9)
  }
}

/**
 * Metadata returned to the client when grounding is enabled.
 *
 * If using grounding with Google Search, you are required to comply with the "Grounding with Google
 * Search" usage requirements for your chosen API provider:
 * [Gemini Developer API](https://ai.google.dev/gemini-api/terms#grounding-with-google-search) or
 * Vertex AI Gemini API (see [Service Terms](https://cloud.google.com/terms/service-terms) section
 * within the Service Specific Terms).
 *
 * @property webSearchQueries The list of web search queries that the model performed to gather the
 * grounding information. These can be used to allow users to explore the search results themselves.
 * @property searchEntryPoint Google Search entry point for web searches. This contains an HTML/CSS
 * snippet that **must** be embedded in an app to display a Google Search Entry point for follow-up
 * web searches related to the model's "Grounded Response".
 * @property groundingChunks The list of [GroundingChunk] classes. Each chunk represents a piece of
 * retrieved content that the model used to ground its response.
 * @property groundingSupports The list of [GroundingSupport] objects. Each object details how
 * specific segments of the model's response are supported by the `groundingChunks`.
 */
public class GroundingMetadata(
  public val webSearchQueries: List<String>,
  public val searchEntryPoint: SearchEntryPoint?,
  public val retrievalQueries: List<String>,
  @Deprecated("Use groundingChunks instead")
  public val groundingAttribution: List<GroundingAttribution>,
  public val groundingChunks: List<GroundingChunk>,
  public val groundingSupports: List<GroundingSupport>,
) {
  @Serializable
  internal data class Internal(
    val webSearchQueries: List<String>?,
    val searchEntryPoint: SearchEntryPoint.Internal?,
    val retrievalQueries: List<String>?,
    @Deprecated("Use groundingChunks instead")
    val groundingAttribution: List<GroundingAttribution.Internal>?,
    val groundingChunks: List<GroundingChunk.Internal>?,
    val groundingSupports: List<GroundingSupport.Internal>?,
  ) {
    internal fun toPublic() =
      GroundingMetadata(
        webSearchQueries = webSearchQueries.orEmpty(),
        searchEntryPoint = searchEntryPoint?.toPublic(),
        retrievalQueries = retrievalQueries.orEmpty(),
        groundingAttribution = groundingAttribution?.map { it.toPublic() }.orEmpty(),
        groundingChunks = groundingChunks?.map { it.toPublic() }.orEmpty(),
        groundingSupports = groundingSupports?.map { it.toPublic() }.orEmpty().filterNotNull()
      )
  }
}

/**
 * Represents a Google Search entry point.
 *
 * @property renderedContent An HTML/CSS snippet that can be embedded in your app. To ensure proper
 * rendering, it's recommended to display this content within a `WebView`.
 * @property sdkBlob A blob of data for the client SDK to render the search entry point.
 */
public class SearchEntryPoint(
  public val renderedContent: String,
  public val sdkBlob: String?,
) {
  @Serializable
  internal data class Internal(
    val renderedContent: String?,
    val sdkBlob: String?,
  ) {
    internal fun toPublic(): SearchEntryPoint {
      // If rendered content is null, the user must not display the grounded result. If they do,
      // they violate the service terms. To prevent this from happening, throw an exception.
      if (renderedContent == null) {
        throw SerializationException("renderedContent is null, should be a string")
      }
      return SearchEntryPoint(renderedContent = renderedContent, sdkBlob = sdkBlob)
    }
  }
}

/**
 * Represents a chunk of retrieved data that supports a claim in the model's response.
 *
 * @property web Contains details if the grounding chunk is from a web source.
 */
public class GroundingChunk(
  public val web: WebGroundingChunk?,
) {
  @Serializable
  internal data class Internal(
    val web: WebGroundingChunk.Internal?,
  ) {
    internal fun toPublic() = GroundingChunk(web = web?.toPublic())
  }
}

/**
 * A grounding chunk from the web.
 *
 * @property uri The URI of the retrieved web page.
 * @property title The title of the retrieved web page.
 * @property domain The domain of the original URI from which the content was retrieved. This is
 * only populated when using the Vertex AI Gemini API.
 */
public class WebGroundingChunk(
  public val uri: String?,
  public val title: String?,
  public val domain: String?
) {
  @Serializable
  internal data class Internal(val uri: String?, val title: String?, val domain: String?) {
    internal fun toPublic() = WebGroundingChunk(uri = uri, title = title, domain = domain)
  }
}

/**
 * Provides information about how a specific segment of the model's response is supported by the
 * retrieved grounding chunks.
 *
 * @property segment Specifies the segment of the model's response content that this grounding
 * support pertains to.
 * @property groundingChunkIndices A list of indices that refer to specific [GroundingChunk] classes
 * within the [GroundingMetadata.groundingChunks] array. These referenced chunks are the sources
 * that support the claim made in the associated `segment` of the response. For example, an array
 * `[1, 3, 4]` means that `groundingChunks[1]`, `groundingChunks[3]`, `groundingChunks[4]` are the
 * retrieved content supporting this part of the response.
 */
public class GroundingSupport(
  public val segment: Segment,
  public val groundingChunkIndices: List<Int>,
) {
  @Serializable
  internal data class Internal(
    val segment: Segment.Internal?,
    val groundingChunkIndices: List<Int>?,
  ) {
    internal fun toPublic(): GroundingSupport? {
      if (segment == null) {
        return null
      }
      return GroundingSupport(
        segment = segment.toPublic(),
        groundingChunkIndices = groundingChunkIndices.orEmpty(),
      )
    }
  }
}

@Deprecated("Use GroundingChunk instead")
public class GroundingAttribution(
  public val segment: Segment,
  public val confidenceScore: Float?,
) {
  @Deprecated("Use GroundingChunk instead")
  @Serializable
  internal data class Internal(
    val segment: Segment.Internal,
    val confidenceScore: Float?,
  ) {
    internal fun toPublic() =
      GroundingAttribution(segment = segment.toPublic(), confidenceScore = confidenceScore)
  }
}

/**
 * Represents a specific segment within a [Content] object, often used to pinpoint the exact
 * location of text or data that grounding information refers to.
 *
 * @property partIndex The zero-based index of the [Part] object within the `parts` array of its
 * parent [Content] object. This identifies which part of the content the segment belongs to.
 * @property startIndex The zero-based start index of the segment within the specified [Part],
 * measured in UTF-8 bytes. This offset is inclusive, starting from 0 at the beginning of the part's
 * content.
 * @property endIndex The zero-based end index of the segment within the specified [Part], measured
 * in UTF-8 bytes. This offset is exclusive, meaning the character at this index is not included in
 * the segment.
 * @property text The text corresponding to the segment from the response.
 */
public class Segment(
  public val startIndex: Int,
  public val endIndex: Int,
  public val partIndex: Int,
  public val text: String,
) {
  @Serializable
  internal data class Internal(
    val startIndex: Int?,
    val endIndex: Int?,
    val partIndex: Int?,
    val text: String?,
  ) {
    internal fun toPublic() =
      Segment(
        startIndex = startIndex ?: 0,
        endIndex = endIndex ?: 0,
        partIndex = partIndex ?: 0,
        text = text ?: ""
      )
  }
}

/**
 * Metadata related to the [UrlContext] tool.
 *
 * @property urlMetadata List of [UrlMetadata] used to provide context to the Gemini model.
 */
@PublicPreviewAPI
public class UrlContextMetadata internal constructor(public val urlMetadata: List<UrlMetadata>) {

  @Serializable
  @PublicPreviewAPI
  internal data class Internal(val urlMetadata: List<UrlMetadata.Internal>?) {
    internal fun toPublic() = UrlContextMetadata(urlMetadata?.map { it.toPublic() } ?: emptyList())
  }
}

/**
 * Metadata for a single URL retrieved by the [UrlContext] tool.
 *
 * @property retrievedUrl The retrieved URL.
 * @property urlRetrievalStatus The status of the URL retrieval.
 */
@PublicPreviewAPI
public class UrlMetadata
internal constructor(
  public val retrievedUrl: String?,
  public val urlRetrievalStatus: UrlRetrievalStatus
) {
  @Serializable
  internal data class Internal(
    val retrievedUrl: String?,
    val urlRetrievalStatus: UrlRetrievalStatus.Internal
  ) {
    internal fun toPublic() = UrlMetadata(retrievedUrl, urlRetrievalStatus.toPublic())
  }
}

/**
 * The status of a URL retrieval.
 *
 * @property name The name of the retrieval status.
 * @property ordinal The ordinal value of the retrieval status.
 */
@PublicPreviewAPI
public class UrlRetrievalStatus
private constructor(public val name: String, public val ordinal: Int) {

  @Serializable(Internal.Serializer::class)
  internal enum class Internal {
    @SerialName("URL_RETRIEVAL_STATUS_UNSPECIFIED") UNSPECIFIED,
    @SerialName("URL_RETRIEVAL_STATUS_SUCCESS") SUCCESS,
    @SerialName("URL_RETRIEVAL_STATUS_ERROR") ERROR,
    @SerialName("URL_RETRIEVAL_STATUS_PAYWALL") PAYWALL,
    @SerialName("URL_RETRIEVAL_STATUS_UNSAFE") UNSAFE;

    internal object Serializer : KSerializer<Internal> by FirstOrdinalSerializer(Internal::class)

    internal fun toPublic() =
      when (this) {
        SUCCESS -> UrlRetrievalStatus.SUCCESS
        ERROR -> UrlRetrievalStatus.ERROR
        PAYWALL -> UrlRetrievalStatus.PAYWALL
        UNSAFE -> UrlRetrievalStatus.UNSAFE
        else -> UrlRetrievalStatus.UNSPECIFIED
      }
  }

  public companion object {
    /** Unspecified retrieval status. */
    @JvmField public val UNSPECIFIED: UrlRetrievalStatus = UrlRetrievalStatus("UNSPECIFIED", 0)

    /** The URL retrieval was successful. */
    @JvmField public val SUCCESS: UrlRetrievalStatus = UrlRetrievalStatus("SUCCESS", 1)

    /** The URL retrieval failed. */
    @JvmField public val ERROR: UrlRetrievalStatus = UrlRetrievalStatus("ERROR", 2)

    /** The URL retrieval failed because the content is behind a paywall. */
    @JvmField public val PAYWALL: UrlRetrievalStatus = UrlRetrievalStatus("PAYWALL", 3)

    /** The URL retrieval failed because the content is unsafe. */
    @JvmField public val UNSAFE: UrlRetrievalStatus = UrlRetrievalStatus("UNSAFE", 4)
  }
}
